import re

class QueryExecutor:
    def __init__(self, query_code):
        self.query_code = self._extract_and_sanitize_code(query_code)
    
    def _extract_and_sanitize_code(self, code):
        # 提取出在 ```python ``` 和 ``` 之间的内容
        code_match = re.search(r"```python(.*?)```", code, re.DOTALL)
        if code_match:
            code = code_match.group(1).strip()
        else:
            # 如果没有找到代码块，假设整个代码字符串都是代码
            code = code.strip()
        
        # 移除行内和块注释
        lines = code.split("\n")
        sanitized_lines = []
        for line in lines:
            # 移除行内注释
            line = re.sub(r"#.*", "", line).strip()
            if line:  # 如果行中还有代码内容，则保留
                sanitized_lines.append(line)
        
        # 将代码行重新组合
        sanitized_code = "\n".join(sanitized_lines)
        return sanitized_code
    
    def execute(self, df):
        local_env = {'df': df, 'result': None}
        
        try:
            exec(self.query_code, {}, local_env)
            result = local_env.get('result')
        except Exception as e:
            result = f"Execution error: {e}"
        
        return result


